/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.mapreduce.benchmark;

import java.util.Arrays;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sss.client.DataPutter;
import org.sss.client.Job;
import org.sss.client.JobEngine;
import org.sss.client.SssClient;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.Output;
import org.sss.mapreduce.Reducer;
import org.sss.mapreduce.datatype.PackableBytes;
import org.sss.mapreduce.datatype.PackableInt;
import org.sss.mapreduce.datatype.PackableLong;

public class WriteBenchmark {
  private static Logger logger = LoggerFactory.getLogger(WriteBenchmark.class);

  public static void main(String[] args) throws Exception {
    SssClient client = new SssClient(args);
    String[] others = client.getArgs();

    if (others.length < 2) {
      System.err.println("org.sss.mapreduce.benchmark.WriteBenchmark TUPLE_SIZE NUM_TUPLES");
      System.exit(1);
    }

    int  recSize = Integer.parseInt(others[0]);
    long nRec    = Long.parseLong(others[1]);

    // upload
    int nNodes = client.getCluster().getStorageNodes().size();
    System.out.println("# each tuple size is " + recSize);
    System.out.println("# number of tuples is " + nRec);
    System.out.println("# Upload data(" + nNodes +  " nodes).");

    PackableBytes content = new PackableBytes();

    DataPutter<PackableInt, PackableBytes> putter =
      DataPutter.create(client, PackableInt.class, PackableBytes.class);
    GroupID input = putter.getOutputGroupID();
    try {
      for (int i = 0;i < nNodes * input.getNSlot();++i) {
        putter.append(new PackableInt(i), content);
      }
    } finally {
      putter.close();
    }

    // Run
    System.out.println("# start");
    long start = System.currentTimeMillis();
    JobEngine engine = new JobEngine(client);
    try {
      GroupID out = GroupID.createRandom(engine);

      Job.Builder jb = engine.getJobBuilder("write-benchmark", WriteBenchmarkReducer.class)
        .addInput(input).addOutput("output", out);
      jb.getConfiguration().setInt("writebenchmark.key.num", nNodes * input.getNSlot());
      jb.getConfiguration().setLong("writebenchmark.tuple.num", nRec);
      jb.getConfiguration().setInt("writebenchmark.tuple.size", recSize);

      jb.build();
      engine.exec();

      long end = System.currentTimeMillis();
      long elapsed = end - start;

      System.out.println("# end. elapsed " + elapsed + " ms.");
      System.out.println(Long.toString(elapsed));
    } finally {
      engine.dispose();
    }
  }

  public static class WriteBenchmarkReducer extends Reducer {
    private long nRec = 0;
    private int nKey = 0;
    private PackableBytes content = null;

    @Override
    public void configure(Context context) throws Exception {
      Configuration conf = context.getConfiguration();
      nKey = conf.getInt("writebenchmark.key.num");
      nRec = conf.getLong("writebenchmark.tuple.num");
      int nSize = conf.getInt("writebenchmark.tuple.size");

      byte[] bs = new byte[nSize];
      Arrays.fill(bs, (byte)0);
      content = new PackableBytes(bs);
    }

    public void reduce(Context context,
                       PackableInt key, Iterable<PackableBytes> values,
                       Output<PackableLong, PackableBytes> output) throws Exception {
      int idx = key.get();

      logger.debug("index = {}, size   = {}.", idx, content.get().length);

      for (long i = idx; i < nRec;i += nKey) {
        output.write(new PackableLong(i), content);
      }
    }
  }
}
